{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b122c690",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import cv2\n",
    "import pickle\n",
    "import supervision as sv\n",
    "from supervision.annotators.utils import ColorLookup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3aaebf4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "info_path = '../data/nuscenes/nuscenes_data_info.pkl'\n",
    "save_dir = '../data/nuscenes/img_results'\n",
    "with open(info_path, 'rb') as f:\n",
    "    data_info = pickle.load(f)\n",
    "classes = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 5: 'bus', 7: 'truck'}\n",
    "\n",
    "cams = data_info[0]['samples'][0]['cams'].keys()\n",
    "scene = data_info[3]\n",
    "scene_name = scene['scene_name']\n",
    "scene_save_dir = f\"{save_dir}/{scene_name}\"\n",
    "# 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT'\n",
    "cam = 'CAM_FRONT'\n",
    "\n",
    "fps = 2\n",
    "w, h = 1600, 900\n",
    "video_writer = cv2.VideoWriter(f\"{scene_name}_{cam}.mp4\", cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (w, h))\n",
    "\n",
    "thickness = 1\n",
    "text_scale = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "32f99305",
   "metadata": {},
   "outputs": [],
   "source": [
    "pseudo_labels_path = '../data/nuscenes/pseudo_labels'\n",
    "with open(f\"{pseudo_labels_path}/{scene_name}.pkl\", 'rb') as f:\n",
    "    pseudo_labels = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "771e89b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in range(len(scene['samples'])):\n",
    "    sample = scene['samples'][idx]\n",
    "    img_path = '.' + sample['cams'][cam]['img_path']\n",
    "    img = cv2.imread(img_path)\n",
    "    save_path = os.path.join(scene_save_dir, cam, os.path.basename(img_path).replace('.jpg', '.pkl'))\n",
    "    img_result = pickle.load(open(save_path, 'rb'))\n",
    "\n",
    "    if len(img_result['bboxes']) == 0:\n",
    "        video_writer.write(img)\n",
    "        continue\n",
    "\n",
    "    bboxes = img_result['bboxes']  # xywh format\n",
    "    labels = img_result['labels']\n",
    "    scores = img_result['scores']\n",
    "    ids = img_result['ids']\n",
    "    masks = img_result['masks']  # list of masks\n",
    "    xyxy = np.concatenate([bboxes[:, :2] - bboxes[:, 2:] / 2, bboxes[:, :2] + bboxes[:, 2:] / 2], axis=1)\n",
    "\n",
    "    img_masks = []\n",
    "    for contour in masks:\n",
    "        img_mask = np.zeros(img.shape[:2], dtype=np.uint8)\n",
    "        cv2.drawContours(img_mask, [contour], -1, 255, thickness=cv2.FILLED)\n",
    "        img_mask = img_mask.astype(bool)\n",
    "        img_masks.append(img_mask)\n",
    "    img_masks = np.array(img_masks)\n",
    "    if ids is not None:\n",
    "        anno_labels = [classes[label] + f'_{id}' for label, id in zip(labels, ids)]\n",
    "    else:\n",
    "        ids = np.arange(len(labels))\n",
    "        anno_labels = [classes[label] for label in labels]\n",
    "\n",
    "    detections = sv.Detections(xyxy=xyxy, confidence=scores, class_id=labels, mask=img_masks, tracker_id=ids)\n",
    "    annotated_image = img.copy()\n",
    "    annotated_image = sv.MaskAnnotator(\n",
    "        color_lookup=sv.ColorLookup.TRACK,\n",
    "        opacity=0.4\n",
    "    ).annotate(scene=annotated_image, detections=detections)\n",
    "    annotated_image = sv.BoxAnnotator(\n",
    "        color_lookup=sv.ColorLookup.TRACK,\n",
    "        thickness=thickness\n",
    "    ).annotate(scene=annotated_image, detections=detections)\n",
    "    annotated_image = sv.LabelAnnotator(\n",
    "        color_lookup=sv.ColorLookup.TRACK,\n",
    "        text_scale=text_scale,\n",
    "        text_padding=0,\n",
    "        smart_position=True\n",
    "    ).annotate(scene=annotated_image, detections=detections, labels=anno_labels)\n",
    "\n",
    "    cam_infos = scene['samples'][idx]['cams']\n",
    "    cam2img = np.eye(4, dtype=np.float32)\n",
    "    cam2img[:3, :3] = cam_infos[cam]['cam2img'][:3, :3]\n",
    "    lidar2img = cam2img @ cam_infos[cam]['lidar2cam']\n",
    "\n",
    "    objects = pseudo_labels[idx]['objects']\n",
    "    bbox_3d_masks = []\n",
    "    for obj in objects:\n",
    "        if obj['cam'] != cam or obj['bbox_3d'] is None:\n",
    "            continue\n",
    "        x, y, z, l, w, h, yaw = obj['bbox_3d']\n",
    "        bbox_3d_corners = np.array([\n",
    "            [l / 2, w / 2, h / 2],\n",
    "            [-l / 2, w / 2, h / 2],\n",
    "            [-l / 2, -w / 2, h / 2],\n",
    "            [l / 2, -w / 2, h / 2],\n",
    "            [l / 2, w / 2, -h / 2],\n",
    "            [-l / 2, w / 2, -h / 2],\n",
    "            [-l / 2, -w / 2, -h / 2],\n",
    "            [l / 2, -w / 2, -h / 2]\n",
    "        ])\n",
    "        rotation_matrix = np.array([\n",
    "            [np.cos(yaw), -np.sin(yaw), 0],\n",
    "            [np.sin(yaw), np.cos(yaw), 0],\n",
    "            [0, 0, 1]\n",
    "        ])\n",
    "        bbox_3d_corners = (rotation_matrix @ bbox_3d_corners.T).T + np.array([x, y, z + h / 2])\n",
    "        corners_homo = np.hstack((bbox_3d_corners, np.ones((bbox_3d_corners.shape[0], 1))))\n",
    "        corners_img = (lidar2img @ corners_homo.T).T\n",
    "        corners_img = (corners_img[:, :2] / np.maximum(corners_img[:, 2:3], 1e-4)).astype(np.float32)\n",
    "        in_img = (corners_img[:, 0] >= 0) & (corners_img[:, 0] < 1600) & (corners_img[:, 1] >= 0) & (corners_img[:, 1] < 900)\n",
    "        if in_img.sum() == 0:\n",
    "            continue\n",
    "        hull = cv2.convexHull(corners_img).astype(np.int32)\n",
    "        bbox_3d_mask = np.zeros(img.shape[:2], dtype=np.uint8)\n",
    "        cv2.fillPoly(bbox_3d_mask, [hull], 255)\n",
    "        bbox_3d_masks.append(bbox_3d_mask)\n",
    "    bbox_3d_masks = np.array(bbox_3d_masks)\n",
    "    # 将bbox_3d_masks应用到annotated_image上，通过增加一个混合mask的方式\n",
    "    for mask in bbox_3d_masks:\n",
    "        annotated_image[mask == 255] = (0, 255, 0)\n",
    "\n",
    "    video_writer.write(annotated_image)\n",
    "\n",
    "video_writer.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b3c0a72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx = 35\n",
    "\n",
    "sample = scene['samples'][idx]\n",
    "img_path = '.' + sample['cams'][cam]['img_path']\n",
    "img = cv2.imread(img_path)\n",
    "save_path = os.path.join(scene_save_dir, cam, os.path.basename(img_path).replace('.jpg', '.pkl'))\n",
    "img_result = pickle.load(open(save_path, 'rb'))\n",
    "\n",
    "bboxes = img_result['bboxes']  # xywh format\n",
    "labels = img_result['labels']\n",
    "scores = img_result['scores']\n",
    "ids = img_result['ids']\n",
    "masks = img_result['masks']  # list of masks\n",
    "xyxy = np.concatenate([bboxes[:, :2] - bboxes[:, 2:] / 2, bboxes[:, :2] + bboxes[:, 2:] / 2], axis=1)\n",
    "\n",
    "img_masks = []\n",
    "for contour in masks:\n",
    "    img_mask = np.zeros(img.shape[:2], dtype=np.uint8)\n",
    "    cv2.drawContours(img_mask, [contour], -1, 255, thickness=cv2.FILLED)\n",
    "    img_mask = img_mask.astype(bool)\n",
    "    img_masks.append(img_mask)\n",
    "img_masks = np.array(img_masks)\n",
    "if ids is not None:\n",
    "    anno_labels = [classes[label] + f'_{id}' for label, id in zip(labels, ids)]\n",
    "else:\n",
    "    ids = np.arange(len(labels))\n",
    "    anno_labels = [classes[label] for label in labels]\n",
    "\n",
    "detections = sv.Detections(xyxy=xyxy, confidence=scores, class_id=labels, mask=img_masks, tracker_id=ids)\n",
    "annotated_image = img.copy()\n",
    "annotated_image = sv.MaskAnnotator(\n",
    "    color_lookup=sv.ColorLookup.TRACK,\n",
    "    opacity=0.4\n",
    ").annotate(scene=annotated_image, detections=detections)\n",
    "annotated_image = sv.BoxAnnotator(\n",
    "    color_lookup=sv.ColorLookup.TRACK,\n",
    "    thickness=thickness\n",
    ").annotate(scene=annotated_image, detections=detections)\n",
    "annotated_image = sv.LabelAnnotator(\n",
    "    color_lookup=sv.ColorLookup.TRACK,\n",
    "    text_scale=text_scale,\n",
    "    text_padding=0,\n",
    "    smart_position=True\n",
    ").annotate(scene=annotated_image, detections=detections, labels=anno_labels)\n",
    "\n",
    "# save\n",
    "save_path = f\"{scene_name}_{cam}_{idx}.jpg\"\n",
    "cv2.imwrite(save_path, annotated_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "532d054f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
